#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
display.py — GPU-accelerated OGG spectrogram visualizer.

Usage:
    python display.py path/to/audio.ogg --scale log
"""

import argparse
from pathlib import Path
import datetime
import torch
import torchaudio
import matplotlib.pyplot as plt
import numpy as np


def compute_spectrogram(audio_path: str, n_fft: int = 1024, hop_length: int = 512) -> tuple[torch.Tensor, int]:
    """
    Compute the spectrogram of an audio file using torchaudio.

    Args:
        audio_path: Path to audio file (ogg/wav/mp3/etc)
        n_fft: FFT window size
        hop_length: Hop length between FFT windows

    Returns:
        (spectrogram, sample_rate)
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Load audio (torchaudio handles OGG automatically)
    waveform, sr = torchaudio.load(audio_path)
    waveform = waveform.to(device)

    # Compute spectrogram
    spec_transform = torchaudio.transforms.Spectrogram(
        n_fft=n_fft,
        hop_length=hop_length,
        power=2.0
    ).to(device)

    spectrogram = spec_transform(waveform)
    if device.type == "cuda":
        print("Processing on GPU completed, moving to CPU...")

    return spectrogram.cpu(), sr


def plot_spectrogram_2d(
    spec: torch.Tensor,
    sr: int,
    hop_length: int,
    title: str = "Spectrogram (2D Brightness Mode)",
    scale: str = "log"
) -> None:
    """
    Plot a 2D spectrogram (time vs frequency, brightness = intensity).

    Args:
        spec: Spectrogram tensor
        sr: Sample rate
        hop_length: Hop length between frames
        title: Plot title
        scale: 'linear' or 'log' frequency axis
    """
    if spec.ndim == 3:
        spec = spec.mean(dim=0)  # mixdown to mono
    elif spec.ndim != 2:
        raise ValueError(f"Unexpected spectrogram shape: {spec.shape}")

    print("Received tensor, rendering...")

    # Convert to decibel scale
    spec_db = 10 * torch.log10(spec + 1e-10)
    spec_np = spec_db.detach().cpu().numpy()

    num_freq_bins, num_frames = spec_np.shape
    time_axis = np.linspace(0, num_frames * hop_length / sr, num_frames)
    freq_axis = np.linspace(0, sr / 2, num_freq_bins)

    # Normalize for brightness
    norm = (spec_np - spec_np.min()) / (spec_np.max() - spec_np.min() + 1e-9)

    # Plot
    plt.figure(figsize=(12, 6))
    plt.imshow(
        norm,
        origin="lower",
        aspect="auto",
        cmap="magma",
        extent=[time_axis[0], time_axis[-1], freq_axis[0], freq_axis[-1]],
    )

    plt.colorbar(label="Normalized Intensity")
    plt.xlabel("Time [s]")
    plt.ylabel("Frequency [Hz]")
    plt.title(title)

    if scale == "log":
        plt.yscale("log")
        plt.ylim(20, sr / 2)  # Avoid log(0)
        plt.ylabel("Frequency [Hz] (log scale)")

    plt.tight_layout()
    out_name = f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{scale}.svg"
    plt.savefig(out_name)
    print(f"Saved to {out_name}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Display audio spectrogram (CUDA-aware)")
    parser.add_argument("audio", type=str, help="Path to audio file (ogg/wav/mp3)")
    parser.add_argument("--n_fft", type=int, default=1024, help="FFT window size")
    parser.add_argument("--hop", type=int, default=512, help="Hop length between frames")
    parser.add_argument("--title", type=str, default="Spectrogram 2D", help="Plot title")
    parser.add_argument("--scale", type=str, choices=["linear", "log"], default="log", help="Frequency scale mode")

    args = parser.parse_args()
    audio_path = Path(args.audio)
    if not audio_path.exists():
        raise FileNotFoundError(f"Audio file not found: {audio_path}")

    spec, sr = compute_spectrogram(str(audio_path), n_fft=args.n_fft, hop_length=args.hop)
    plot_spectrogram_2d(spec, sr, hop_length=args.hop, title=args.title, scale=args.scale)


if __name__ == "__main__":
    main()
